#ifndef ITER_SOLVER_HPP
#define ITER_SOLVER_HPP

#include "precond.hpp"
#include "utils.hpp"

/// @brief 迭代求解器的基类
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam pc_data_t 预条件子存储类型，可以为float或者double
/// @tparam pc_calc_t 预条件子计算类型，可以为float或者double
/// @tparam setup_t 外迭代类型，可以为float或者double
template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t>
class IterSolver
{
public:
    bool inited = false; // 是否已经初始化
    Precond<idx_t, pc_data_t, pc_calc_t, ksp_t>* precond = nullptr; // 预条件子
    idx_t max_iter = 10; // 最大迭代次数
    ksp_t rel_tol = 0.0, abs_tol = 0.0;// 用高精度的收敛判断
    mutable idx_t final_iter = 0, converged = 0; // 最终迭代次数，是否收敛
    mutable ksp_t final_norm = 0.0; // 最终残差范数
    mutable ksp_t t_prec = 0.0, t_spmv = 0.0; // 预条件子和矩阵向量乘的时间
    MyParCSRMatrix<idx_t, ksp_t, ksp_t> * oper; // 系数矩阵

public:
    /// @brief 迭代求解器的构造函数
    IterSolver() {}
    /// @brief 迭代求解器的析构函数
    ~IterSolver() {}
    /// @brief 设置预条件子
    /// @param precond 预条件子
    void SetPrecond(Precond<idx_t, pc_data_t, pc_calc_t, ksp_t> *precond) { this->precond = precond; }
    /// @brief 设置收敛判断的相对精度
    /// @param rtol 相对精度
    void SetRelTol(ksp_t rtol) { rel_tol = rtol; }
    /// @brief 设置收敛判断的绝对精度
    /// @param atol 绝对精度
    void SetAbsTol(ksp_t atol) { abs_tol = atol; }
    /// @brief 设置最大迭代次数
    /// @param max_it 最大迭代次数
    void SetMaxIter(idx_t max_it) { max_iter = max_it; }
    /// @brief 获取最终迭代次数
    /// @return 最终迭代次数
    idx_t GetNumIterations() const { return final_iter; }
    /// @brief 获取是否收敛
    /// @return 是否收敛
    idx_t GetConverged() const { return converged; }
    /// @brief 获取最终残差范数
    /// @return 最终残差范数
    ksp_t GetFinalNorm() const { return final_norm; }
    /// @brief 获取预条件子的时间
    /// @return 预条件子的时间
    ksp_t GetPrecTime() const { return t_prec; }
    /// @brief 获取矩阵向量乘的时间
    /// @return 矩阵向量乘的时间
    ksp_t GetSpmvTime() const { return t_spmv; }

    /// @brief 迭代求解器的建立函数
    /// @param A 系数矩阵
    /// @note 1. 保存系数矩阵的指针 2. 预条件子Setup
    virtual void Setup(MyParCSRMatrix<idx_t, ksp_t, ksp_t> & A) {
        if (this->precond) {
            this->precond->Setup(A);
        }
        this->oper = & A;
    }
    /// @brief 迭代求解器的求解函数
    /// @param b 右端项
    /// @param x 解向量
    virtual void Solve(const MyParCSRVector<idx_t, pc_calc_t> & b, MyParCSRVector<idx_t, pc_calc_t> & x) = 0;
    
    /// @brief 矩阵向量乘
    /// @param x 源向量
    /// @param y 目标向量
    /// @note y = A * x
    void Spmv(const MyParCSRVector<idx_t, pc_calc_t> & x, MyParCSRVector<idx_t, pc_calc_t> & y) {
        this->oper->Mult(1.0, x, 0.0, y);
    }
    /// @brief 计算残差
    /// @param x 解向量
    /// @param b 右端项
    /// @param y 目标向量
    /// @note y = b - A * x
    void CalcResi(const MyParCSRVector<idx_t, pc_calc_t> & x, const MyParCSRVector<idx_t, pc_calc_t> & b, MyParCSRVector<idx_t, pc_calc_t> & y) {
        this->oper->Mult(-1.0, x, 1.0, b, y);
    }
};

/// @brief GMRES 迭代求解器
template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t>
class GMRESSolver: public IterSolver<idx_t, pc_data_t, pc_calc_t, ksp_t>
{
public:
    int restart_len = 10; // 重启长度
public:
    /// @brief 设置重启长度
    /// @param num 重启长度
    virtual void SetRestartlen(int num) { restart_len = num; }
    /// @brief GMRES 迭代求解器的求解函数
    /// @param b 右端项
    /// @param x 解向量
    virtual void Solve(const MyParCSRVector<idx_t, pc_calc_t> & b, MyParCSRVector<idx_t, pc_calc_t> & x);
};

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t>
void GMRESSolver<idx_t, pc_data_t, pc_calc_t, ksp_t>::Solve(const MyParCSRVector<idx_t, pc_calc_t> & b, MyParCSRVector<idx_t, pc_calc_t> & x) {
    int my_pid; 
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    double & t_prec = this->t_prec = 0.0;
    double & t_spmv = this->t_spmv = 0.0;

    // 分配辅助空间
    ksp_t * rs = new ksp_t [restart_len + 1];
    ksp_t *  c = new ksp_t [restart_len];
    ksp_t *  s = new ksp_t [restart_len];
    ksp_t * rs2= new ksp_t [restart_len + 1];
    ksp_t ** hh = new ksp_t * [restart_len + 1];
    for (int i = 0; i <= restart_len; i++)
        hh[i] = new ksp_t [restart_len];

    MyParCSRVector<idx_t, ksp_t> r(x), w(x);
    std::allocator<MyParCSRVector<idx_t, ksp_t> > alloc;
    MyParCSRVector<idx_t, ksp_t> * p = alloc.allocate(restart_len + 1);
    for (int i = 0; i <= restart_len; i++) {
        alloc.construct(p + i, x);
    }

    // 低精度(若有)向量缓冲区
    pc_calc_t * pc_buf_b = nullptr, * pc_buf_x = nullptr;
    if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
        pc_buf_b = new MyParCSRVector<idx_t, pc_calc_t>(x);
        pc_buf_x = new MyParCSRVector<idx_t, pc_calc_t>(x);
    }

    ksp_t b_norm = vec_norm(b);
    ksp_t r_norm, real_r_norm_old = b_norm, real_r_norm_new;

    // 计算初始残差
    t_spmv -= MPI_Wtime();
    this->CalcResi(x, b, p[0]);
    t_spmv += MPI_Wtime();

    r_norm = vec_norm(p[0]);

    if (my_pid == 0) {
        printf("L2 norm of b: %20.16e\n", (double)b_norm);
        printf("Initial L2 norm of residual: %20.16e\n", (double)r_norm);
    }

    int & iter = this->final_iter = 0;
    ksp_t den_norm;
    if (b_norm > 0.0)   den_norm = b_norm;// convergence criterion |r_i|/|b| <= accuracy if |b| > 0
    else                den_norm = r_norm;// convergence criterion |r_i|/|r0| <= accuracy if |b| = 0

    /*  convergence criteria: |r_i| <= max( a_tol, r_tol * den_norm)
        den_norm = |r_0| or |b|
        note: default for a_tol is 0.0, so relative residual criteria is used unless
        user specifies a_tol, or sets r_tol = 0.0, which means absolute
        tol only is checked  */
    ksp_t epsilon;
    epsilon = std::max(this->abs_tol, this->rel_tol * den_norm);

    // 收敛判据：|r_i| <= epsilon
    if (my_pid == 0) {
        if (b_norm > 0.0) {
            printf("=========================================================\n\n");
            printf("Iters             resid.norm               rel.res.norm  \n");
            printf("------       -------------------      --------------------\n");
        } else {
            printf("=============================================\n\n");
            printf("Iters     resid.norm     \n");
            printf("------    ------------    \n");
        }
    }

    // 开始迭代
    while (iter < this->max_iter) {
        // 初始化 hessenberg 系统的第一项
        rs[0] = r_norm;
        if (r_norm == 0.0) {
            this->converged = 1;
            goto finish;
        }

        // 查看是否已经收敛
        if (r_norm <= epsilon) {
            t_spmv -= MPI_Wtime();
            this->CalcResi(x, b, r);
            t_spmv += MPI_Wtime();

            r_norm = vec_norm(r);
            if (r_norm <= epsilon) {
                if (my_pid == 0)
                    printf("\nFinal L2 norm of residual: %20.16e\n\n", (double)r_norm);
                break;
            } else {
                if (my_pid == 0) 
                    printf("false convergence\n");
            }
        }

        // 初始化 p[0] = r / r_norm
        vec_scale(p[0], 1.0 / r_norm, p[0]);
        int i = 0;

        // 重启循环
        while (i < restart_len && iter < this->max_iter) {
            i++;
            iter++;

            // 预条件子
            t_prec -= MPI_Wtime();
            if (this->precond) {
                if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                    // 混合精度需要额外转换向量精度
                    vec_copy(*pc_buf_b, p[i-1]);
                    vec_copy(*pc_buf_x, r);
                    this->precond->Solve(*pc_buf_b, *pc_buf_x);
                    vec_copy(r, *pc_buf_x);
                } else {
                    this->precond->Solve(p[i-1], r);
                }
            }
            else vec_copy(r, p[i-1]); // 无预条件子
            t_prec += MPI_Wtime();

            // p[i] = A * r
            t_spmv -= MPI_Wtime();
            this->Spmv(r, p[i]);
            t_spmv += MPI_Wtime();

            // modified Gram_schmidt
            for (int j = 0; j < i; j++) {
                hh[j][i-1] = vec_dot(p[j], p[i]);
                vec_add(p[i], -hh[j][i-1], p[j], p[i]);
            }

            // Arnoldi step
            ksp_t t = vec_norm(p[i]);
            hh[i][i-1] = t;	
            if (t != 0.0) vec_scale(p[i], 1.0 / t, p[i]);

            // 更新 hh 矩阵
            for (int j = 1; j < i; j++) {
                ksp_t t = hh[j-1][i-1];
                hh[j-1][i-1] = s[j-1]*hh[j][i-1] + c[j-1] * t;
                hh[j][i-1]   = -s[j-1]*t + c[j-1]*hh[j][i-1];
            }
            t = hh[i][i-1] * hh[i][i-1];
            t += hh[i-1][i-1] * hh[i-1][i-1];
            ksp_t gamma = sqrt(t);
            if (gamma == 0.0) gamma = 1.0e-16;
            c[i-1] = hh[i-1][i-1]/gamma;
            s[i-1] = hh[i][i-1]/gamma;
            rs[i] = -hh[i][i-1]*rs[i-1];
            rs[i]/=  gamma;
            rs[i-1] = c[i-1]*rs[i-1];

            // 得出残差范数
            hh[i-1][i-1] = s[i-1]*hh[i][i-1] + c[i-1]*hh[i-1][i-1];
            r_norm = fabs(rs[i]);

            if (my_pid == 0) {
                if (b_norm > 0.0)
                    printf("%5d    %.16e   %.16e\n", iter, (double)r_norm, (double)(r_norm/b_norm));
                else
                    printf("%5d    %.16e\n", iter, (double)r_norm);
            }

            if (r_norm <= epsilon) break;
        } // end of restart while loop

        // 计算解向量
        // 求解上三角系统：hh * y = rs
        rs[i-1] = rs[i-1]/hh[i-1][i-1];
        for (int k = i-2; k >= 0; k--)
        {
            ksp_t t = 0.0;
            for (int j = k+1; j < i; j++) {
                t -= hh[k][j]*rs[j];
            }
            t += rs[k];
            rs[k] = t/hh[k][k];
        }

        // 计算解向量
        vec_scale(p[i-1], rs[i-1], w);
        for (int j = i-2; j >= 0; j--)
            vec_add(w, rs[j], p[j], w);

        // 预条件子
        t_prec -= MPI_Wtime();
        if (this->precond) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                vec_copy(*pc_buf_b, w);
                vec_copy(*pc_buf_x, r);
                this->precond->Solve(*pc_buf_b, *pc_buf_x);
                vec_copy(r, *pc_buf_x);
            } else {
                this->precond->Solve(w, r);
            }
        }
        else vec_copy(r, w);
        t_prec += MPI_Wtime();

        // 更新解向量
        vec_add(x, 1.0, r, x);

        // 计算真实残差
        if (r_norm <= epsilon) {
            t_spmv -= MPI_Wtime();
            this->CalcResi(x, b, r);
            t_spmv += MPI_Wtime();

            real_r_norm_new = r_norm = vec_norm(r);

            if (r_norm <= epsilon) {
                if (my_pid == 0)
                    printf("\nFinal L2 norm of residual: %20.16e\n\n", (double)r_norm);
                goto finish;
            } else { // 真实残差尚未收敛
                // 若真实残差没有减小，则退出
                if (real_r_norm_new >= real_r_norm_old) {
                    if (my_pid == 0)
                        printf("\nFinal L2 norm of residual: %20.16e\n\n", (double)r_norm);
                    this->converged = 1;
                    break;
                }

                // 报告真实残差和重启残差的差异
                if (my_pid == 0)
                    printf("false convergence 2, L2 norm of residual: %20.16e\n", (double)r_norm);

                vec_copy(p[0], r);
                i = 0;
                real_r_norm_old = real_r_norm_new;
            }
        }

        // 计算残差向量，继续迭代
	    for (int j=i ; j > 0; j--) {
            rs[j-1] = -s[j-1]*rs[j];
            rs[j] = c[j-1]*rs[j];
    	}
        
        if (i) vec_add(p[i], rs[i] - 1.0, p[i], p[i]);

        for (int j=i-1 ; j > 0; j--)
            vec_add(p[i], rs[j], p[j], p[i]);
        
        if (i) {
            vec_add(p[0], rs[0] - 1.0, p[0], p[0]);
            vec_add(p[0], 1.0, p[i], p[0]);
        }

        if (my_pid == 0) printf("Restarting...\n");
    }/* END of iteration while loop */

finish:
    delete []c; delete []s; delete []rs; delete []rs2;
    c = s = rs = rs2 = nullptr;
    for (int i = 0; i <= restart_len; i++) {
        delete hh[i]; hh[i] = nullptr; 
    }
    delete []hh; hh = nullptr;
    for (int i = 0; i <= restart_len; i++) alloc.destroy(p + i);
    alloc.deallocate(p, restart_len + 1);
    if (pc_buf_x) delete pc_buf_x;
    if (pc_buf_b) delete pc_buf_b;
}

#endif // ITER_SOLVER_HPP